# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags
from coin.flow.topic_record import (TopicRecordReader,
                                    TopicRecordWriter,
                                    TopicRecordCompressionType,
                                    TopicRecord)

FLAGS = flags.FLAGS


def main(_):
  root_dir = flags.FLAGS.root_dir
  out_root_dir = flags.FLAGS.out_root_dir
  # input_file_name = flags.FLAGS.input_file_name
  # output_file_name = flags.FLAGS.output_file_name
  # topic_id = flags.FLAGS.topic_id

  assert os.path.isdir(root_dir)
  for file_name in os.listdir(root_dir):
    if file_name.startswith('huobi_ws_depth'):
      modified_file_name = file_name.replace('depth', 'trade')
      topic_id = 0x0000ED0F000A1200
    elif file_name.startswith('huobi_ws_trade'):
      modified_file_name = file_name.replace('trade', 'depth')
      topic_id = 0x0000ED0F000A1100
    elif file_name.startswith('huobi_ws_futures_depth'):
      modified_file_name = file_name.replace('depth', 'trade')
      topic_id = 0x0000ED0F000A1400
    elif file_name.startswith('huobi_ws_futures_trade'):
      modified_file_name = file_name.replace('trade', 'depth')
      topic_id = 0x0000ED0F000A1300
    else:
      continue
    reader = TopicRecordReader(os.path.join(root_dir, file_name),
                               compression_type=TopicRecordCompressionType.GZIP)

    writer = TopicRecordWriter(os.path.join(out_root_dir, modified_file_name),
                               compression_type=TopicRecordCompressionType.GZIP)

    for record in reader.read_records():
      modified_record = TopicRecord(queue_seq=record.queue_seq,
                                    topic_id=topic_id,
                                    topic_seq=record.topic_seq,
                                    topic_str=None,
                                    timestamp=record.timestamp,
                                    data=record.data)
      writer.write_record(modified_record)

    writer.flush()
    writer.close()


if __name__ == '__main__':
  flags.DEFINE_string('root_dir', None, 'Root dir.')

  flags.DEFINE_string('out_root_dir', None, 'Out root dir.')

  flags.DEFINE_string('input_file_name', None, 'Input.')

  flags.DEFINE_string('output_file_name', None, 'Output.')

  flags.DEFINE_integer('topic_id', None, 'topic_id.')

  app.run(main)
